#!/usr/bin/env python
"""
PyLink-SR (Python Linked Software Radio System)
(c) Bob Cowdery (G3UKB) 2008

File: dsp.py
Date: July 2008

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

The author can be reached by email at

bob@g3ukb.co.uk

========================================================================
Module Description:

The signal processing

"""

#-----------------------------------------------------------------------
# System imports
import logging
import sys
from time import sleep
from scipy import *

# Application imports
#from defs import *
#import broker
import config_service as config

#-----------------------------------------------------------------------
# The single DSP class
class Dsp:
    
    #-------------------------------------------------------------------
    def __init__(self, event_service):
        
        # Stash the event service
        self.es = event_service
        
        # declare the working arrays
        FRAMES = config.frames()
        self.complex_smpls = array([])
        self.complex_fd = array([])
        self.processed_smpls = array([zeros(FRAMES,'d'),zeros(FRAMES,'d')])
        self.overlap = zeros(FRAMES*2,'D')
        
        # Set some defaults
        self.current_filter = '2K4'
        self.current_mode = 'LSB'
        self.current_agc = True
        self.set_mode_filter('LSB', '2K4')
        self.current_agc_rate = 'AGC_FAST'
        
        # AGC variables
        AGC_GAIN_LEN = 500                            # agc gain history buffer max size
        self.agc_loop = 0                             # loop counter for agc processing
        self.agc_gain_buf = zeros(AGC_GAIN_LEN,'d')   # agc gain buffer
        self.previous_agc_gain = 0.0                  # track agc gain slope
        # AGC configuration values
        self.ramp = config.dsp_agc_ramp()
        self.gain = config.dsp_agc_gain()
        self.roof = config.dsp_agc_roof()
    
    #-----------------------------------------------------------------------
    # Message handlers
    def set_mode(self, data):  
        
        # Calculate the new filter on a mode change
        self.current_mode = data
        self.set_mode_filter(self.current_mode, self.current_filter)
        self.es.publish('ui_hints', ('mode_hints', data))
        
    #-----------------------------------------------------------------------
    def set_filter(self, data): 
        
        # Calculate the new filter on a filter change
        self.current_filter = data
        self.set_mode_filter(self.current_mode, self.current_filter)
        self.es.publish('ui_hints', ('filter_hints', data))
    
    #-----------------------------------------------------------------------
    def set_agc(self, data): 
        
        # Set the new AGC state
        self.current_agc = data
        self.es.publish('ui_hints', ('agc_hints', ('state', data)))

    #-----------------------------------------------------------------------
    def set_agc_rate(self, data): 
        
        # Set the new AGC rate
        self.current_agc_rate = data
        self.es.publish('ui_hints', ('agc_hints', ('rate', data)))
        
    #===========================================================================
    # Signal Processing
    #-----------------------------------------------------------------------    
    # Main RX Path
    def do_rx(self, raw_samples):
        
        # Samples arrive in a 2d matrix
        sampl_sz = len(raw_samples[0])
        
        # Get into complex notation 
        # Get the I & Q the right way round, not sure how to tell this, apart from listening
        # and ensuring LSB and USB are not reversed
        self.complex_smpls = raw_samples[1][:sampl_sz] + 1j*raw_samples[0][:sampl_sz]
        
        # zero stuff to double length
        self.complex_smpls = concatenate((self.complex_smpls,zeros(sampl_sz,'D')))
        
        # perform fft
        # number of bins will be = to the number of samples
        self.complex_fd = fft(self.complex_smpls)
        
        # filter and demodulate
        self.complex_fd = self.fc_filter * self.complex_fd
        
        # inverse fft
        complexout = ifft(self.complex_fd)
        
        # overlap add
        complexout = complexout + self.overlap
        self.overlap = concatenate((complexout[sampl_sz:sampl_sz*2], zeros(sampl_sz,'D')))    
        
        # apply agc
        complexout = self.blk_digital_agc(complexout, sampl_sz, self.current_agc, self.current_agc_rate, self.ramp, self.gain, self.roof)
        
        # generate output samples
        if self.current_mode == 'AM':
            # Copy the magnitude to both channels
            self.processed_smpls[0] = abs(complexout)[:sampl_sz]
            self.processed_smpls[1] = abs(complexout)[:sampl_sz]
        else:
            # Copy real part of output to both channels
            self.processed_smpls[0] = real(complexout)[:sampl_sz] 
            self.processed_smpls[1] = real(complexout)[:sampl_sz] 
        
        return self.processed_smpls
    
    #===========================================================================
    # Private
    #---------------------------------------------------------------------------       
    # Filter Related
    def blk_filter_calc_ssb(self, tapsz, rate, flow, fhigh, mode):
        
        """
        Calculate complex fast convolution filter for demodulation and filtering of SSB.
        This filter will down convert by 1/4 of the sample rate.
        A rate of 48000 therefore ends up with an IF of 12000.
        The opposite sideband will be suppressed.
    
        Inputs:
          tapsz         -- number of taps in filter (determines skirt)
          rate          -- sample rate
          flow          -- low cutoff
          fhigh         -- high cutoff 
          sideband      -- LSB='L', USB='U'
    
        Outputs:
          filter      -- complex frequency domain response 
    
        """
        FRAMES = config.frames()
        # calculate the high/low cutoff freq as a fraction of the sample rate
        fhc = fhigh/rate
        flc = flow/rate
    
        # Calculate FIR lowpass filter coefficients with n taps
        # This has a cutoff of (fHigh-fLow)/2, so for a 2400Hz filter it is 1200Hz
        # In the complex domain this is +-1200Hz = 2400Hz
        # Note that firwin frequency is relative to Nyquist frequency
        coefs = signal.firwin(tapsz,(fhc-flc))
        
        # Modulate to USB centre frequency
        # we effectively mix by a complex sinusoid at 12000 + filter centre (flc+fhc)
        # phase_i = i*(2*pi*(12000+1500)/48000) if our sample rate is 48000 and filter 2400Hz, centre 1500Hz. 
        # i=0,1,2,...,n 
        # Si = (cos(phase_i) + j*sin(phase_i)) 
        filterFreq = (flc+fhc)*pi
        # create array range 0 to no. of taps
        i = arange(0,tapsz)
        # create the complex filter by mixing the filter and the sinusoid using the identity exp(jwt)=coswt+jsinwt
        self.fc_filter = coefs*exp(1j*filterFreq*i)
                       
        # for LSB create conjugate of USB filter
        # flip to the negative freq side for LSB
        if (mode == 'LSB') or (mode == 'CWL'):
            self.fc_filter = conj(self.fc_filter)
    
        # make the array the same number of samples as the input data
        zerostuff = zeros((FRAMES - tapsz) + FRAMES,'D')
        # add the complement of zeros to match the array length
        self.fc_filter = concatenate((self.fc_filter,zerostuff))
                
        # compute complex frequency domain response of LSB or USB filter
        self.fc_filter = fft(self.fc_filter)
    
    #---------------------------------------------------------------------------
    def blk_filter_calc_ampl(self, tapsz, rate, flow, fhigh, mode ):
           
        """
        Calculate complex fast convolution filter for demodulation and filtering of AM/SAM.
        See blk_filter_calc_ssb for comments.
        Inputs:
          tapsz     -- number of taps in filter (determines skirt)
          rate      -- sample rate
          fhigh     -- high cutoff 
        Outputs:
          m_filter  -- complex frequency domain response 
        """
        FRAMES = config.frames()    
        # calculate the high cutoff freq as a fraction of the sample rate
        fhc = fhigh/rate
        # calculate FIR lowpass filter coefficients with n taps
        coefs = signal.firwin(tapsz, fhc*2)
        # make the array the same number of samples as the input data
        zerostuff = zeros((FRAMES - tapsz) + FRAMES,'D')
        # add the complement of zeros to match the array length
        self.fc_filter = concatenate((coefs, zerostuff))
                
        # compute complex frequency domain response of AM filter with 4096 bins
        self.fc_filter = fft(self.fc_filter)
            
    #-----------------------------------------------------------------------
    def set_mode_filter(self, mode, filter): 
        TAPSZ = config.frames()
        RATE = config.smpl_rate()
        
        if mode == 'AM':
            proc = self.blk_filter_calc_ampl
        else:
            proc = self.blk_filter_calc_ssb
            
        if filter == '6K0':
            proc( TAPSZ, float(RATE), 100.0, 6100.0, mode )
        elif filter == '2K4':
            proc( TAPSZ, float(RATE), 300.0, 2700.0, mode )
        elif filter == '2K1':
            proc( TAPSZ, float(RATE), 300.0, 2400.0, mode )
        elif filter == '1K0':
            proc( TAPSZ, float(RATE), 300.0, 1300.0, mode )
        elif filter == '500':
            proc( TAPSZ, float(RATE), 500.0, 1000.0, mode )
        elif filter == '250':
            proc( TAPSZ, float(RATE), 600.0, 850.0, mode )
        elif filter == '100':
            proc( TAPSZ, float(RATE), 700.0, 800.0, mode )
        else:
            proc( TAPSZ, float(RATE), 300.0, 2700.0, mode )
                    
    #---------------------------------------------------------------------------
    # DSP function blocks
    def blk_digital_agc(self, complex_samples, sampl_sz, agcstate, rate, ramp, gain, roof):
        """
        Apply digital agc with configurable state, ramp period and hang.
        This is a single time constant process.
        
        Inputs:
            agcstate        -- True if AGC on
            rate            -- period to hang before starting to ramp signal
            ramp            -- period in ms over which to ramp up/down gain
            gain            -- gain value
            roof            -- maximum agc gain           
            complex_samples -- processed samples
        Outputs:
            complex_samples -- gain adjusted processed samples 
        """
        hang = 3
        # Play with these numbers so it sounds right
        if rate == 'AGC_FAST': hang = 3
        if rate == 'AGC_MED': hang = 7
        if rate == 'AGC_SLOW': hang = 12

        if agcstate == True:
            # AGC is on
            # Cycle round hang value
            if self.agc_loop < hang-1:
                self.agc_loop = self.agc_loop+1
            else:
                self.agc_loop = 0
    
            # Find peak gain reduction
            # Get peak magnitude from samples
            peak = abs(complex_samples[argmax(abs(complex_samples[:sampl_sz]))])
            if peak != 0:
                self.agc_gain_buf[self.agc_loop] = gain/peak
                # Get the minimum value from the first 'hang' elements of m_agcGainBuf
                mingain = self.agc_gain_buf[argmin(self.agc_gain_buf[:hang])]
            else:
                # Silence input so skip for now
                return complex_samples
            
            # Limit to max allowed gain
            if mingain > roof: 
                mingain = roof
            # Ramp in the appropriate direction
            i=arange(0,ramp)
            if mingain < self.previous_agc_gain:
                # Gain is decreasing, use 1ms attack time
                # Ramp gain down over period (nominal 48 samples = 1ms)
                # The ramp is correct abut should be derived from sample rate 48000*0.001 = 48
                gainstep = (self.previous_agc_gain-mingain)/ramp
                complex_samples[:ramp] = complex_samples[:ramp]*(self.previous_agc_gain-((i+1)*gainstep))
                complex_samples[ramp:sampl_sz] = complex_samples[ramp:sampl_sz]*mingain
            elif mingain > self.previous_agc_gain:
                # Gain is increasing, use 1ms attack time
                # Ramp gain up over 1ms period
                gainstep = (mingain-self.previous_agc_gain)/ramp
                complex_samples[:ramp] = complex_samples[:ramp]*(self.previous_agc_gain+((i+1)*gainstep))
                complex_samples[ramp:sampl_sz] = complex_samples[ramp:sampl_sz]*mingain
            else:
                # Multiply the envelope by the agc gain
                complex_samples = complex_samples*mingain
            # Remember previous gain
            self.previous_agc_gain = mingain
            
        return complex_samples
